use crate::errors::ExecutionError;
use crate::{Dag, EdgeHavePorts, NodeKind};

use crate::node::{OutputPortType, PortHandle};
use daggy::petgraph::graph::EdgeReference;
use daggy::petgraph::visit::{EdgeRef, IntoEdges, IntoEdgesDirected, IntoNodeReferences, Topo};
use daggy::petgraph::Direction;
use daggy::{NodeIndex, Walker};
use dozer_types::log::{error, info};
use dozer_types::node::NodeHandle;
use dozer_types::serde::{Deserialize, Serialize};
use dozer_types::types::Schema;
use std::collections::{HashMap, HashSet};
use std::fmt::Debug;

use super::node::OutputPortDef;
use super::{EdgeType as DagEdgeType, NodeType};

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(crate = "dozer_types::serde")]
pub struct EdgeType {
    pub output_port: PortHandle,
    pub input_port: PortHandle,
    pub schema: Schema,
    pub edge_kind: EdgeKind,
}

impl EdgeType {
    pub fn new(
        output_port: PortHandle,
        input_port: PortHandle,
        schema: Schema,
        edge_kind: EdgeKind,
    ) -> Self {
        Self {
            output_port,
            input_port,
            schema,
            edge_kind,
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(crate = "dozer_types::serde")]
pub enum EdgeKind {
    FromSource {
        port_type: OutputPortType,
        port_name: String,
    },
    FromProcessor,
}

pub trait EdgeHaveSchema: EdgeHavePorts {
    fn schema(&self) -> &Schema;
}

impl EdgeHavePorts for EdgeType {
    fn output_port(&self) -> PortHandle {
        self.output_port
    }

    fn input_port(&self) -> PortHandle {
        self.input_port
    }
}

impl EdgeHaveSchema for EdgeType {
    fn schema(&self) -> &Schema {
        &self.schema
    }
}

#[derive(Debug)]
/// `DagSchemas` is a `Dag` with validated schema on the edge.
pub struct DagSchemas {
    graph: daggy::Dag<NodeType, EdgeType>,
}

impl DagSchemas {
    pub fn into_graph(self) -> daggy::Dag<NodeType, EdgeType> {
        self.graph
    }

    pub fn graph(&self) -> &daggy::Dag<NodeType, EdgeType> {
        &self.graph
    }
}

impl DagSchemas {
    /// Validate and populate the schemas, the resultant DAG will have the exact same structure as the input DAG,
    /// with validated schema information on the edges.
    pub async fn new(dag: Dag) -> Result<Self, ExecutionError> {
        validate_connectivity(&dag);

        match populate_schemas(dag.into_graph()).await {
            Ok(graph) => {
                info!("[pipeline] Validation completed");
                Ok(Self { graph })
            }
            Err(e) => {
                error!("[pipeline] Validation error: {}", e);
                Err(e)
            }
        }
    }

    pub fn collect_ancestor_sources(&self, node_index: NodeIndex) -> HashSet<NodeHandle> {
        let mut sources = HashSet::new();
        collect_ancestor_sources_recursive(self, node_index, &mut sources);
        sources
    }
}

fn collect_ancestor_sources_recursive(
    dag: &DagSchemas,
    node_index: NodeIndex,
    sources: &mut HashSet<NodeHandle>,
) {
    for edge in dag.graph().edges_directed(node_index, Direction::Incoming) {
        let source_node_index = edge.source();
        let source_node = &dag.graph()[source_node_index];
        if matches!(source_node.kind, NodeKind::Source(_)) {
            sources.insert(source_node.handle.clone());
        }
        collect_ancestor_sources_recursive(dag, source_node_index, sources);
    }
}

pub trait DagHaveSchemas {
    type NodeType;
    type EdgeType: EdgeHaveSchema;

    fn graph(&self) -> &daggy::Dag<Self::NodeType, Self::EdgeType>;

    fn get_node_input_schemas(&self, node_index: NodeIndex) -> HashMap<PortHandle, Schema> {
        let mut schemas = HashMap::new();

        for edge in self.graph().edges_directed(node_index, Direction::Incoming) {
            let edge = edge.weight();
            let schema = edge.schema();
            schemas.insert(edge.input_port(), schema.clone());
        }

        schemas
    }

    fn get_node_output_schemas(&self, node_index: NodeIndex) -> HashMap<PortHandle, Schema> {
        let mut schemas = HashMap::new();

        for edge in self.graph().edges(node_index) {
            let edge = edge.weight();
            let schema = edge.schema();
            schemas.insert(edge.output_port(), schema.clone());
        }

        schemas
    }
}

impl DagHaveSchemas for DagSchemas {
    type NodeType = NodeType;
    type EdgeType = EdgeType;

    fn graph(&self) -> &daggy::Dag<Self::NodeType, Self::EdgeType> {
        &self.graph
    }
}

fn validate_connectivity(dag: &Dag) {
    // Every source or processor has at least one outgoing edge.
    for (node_index, node) in dag.graph().node_references() {
        match &node.kind {
            NodeKind::Source(_) | NodeKind::Processor(_) => {
                if dag.graph().edges(node_index).count() == 0 {
                    panic!("Node {} has no outgoing edge", node.handle);
                }
            }
            NodeKind::Sink(_) => {}
        }
    }

    // Processor and sink has at least one input port. Every input port has exactly one incoming edge.
    for (node_index, node) in dag.graph().node_references() {
        let mut input_ports = match &node.kind {
            NodeKind::Source(_) => continue,
            NodeKind::Processor(processor) => processor.get_input_ports(),
            NodeKind::Sink(sink) => sink.get_input_ports(),
        };
        if input_ports.is_empty() {
            panic!("Node {} has no input port", node.handle);
        }

        input_ports.sort();

        let mut connected_input_ports = dag
            .graph()
            .edges_directed(node_index, Direction::Incoming)
            .map(|edge| edge.weight().to)
            .collect::<Vec<_>>();
        connected_input_ports.sort();

        if input_ports != connected_input_ports {
            panic!(
                "Node {} has input ports {input_ports:?}, but the incoming edges are {connected_input_ports:?}",
                node.handle
            );
        }
    }
}

/// In topological order, pass output schemas to downstream nodes' input schemas.
async fn populate_schemas(
    dag: daggy::Dag<NodeType, DagEdgeType>,
) -> Result<daggy::Dag<NodeType, EdgeType>, ExecutionError> {
    let mut edges = vec![None; dag.graph().edge_count()];

    for node_index in Topo::new(&dag).iter(&dag) {
        let node = &dag.graph()[node_index];

        match &node.kind {
            NodeKind::Source(source) => {
                let ports = source.get_output_ports();

                for edge in dag.graph().edges(node_index) {
                    let port = edge.weight().from;
                    let port_type = find_output_port_type(&ports, edge);
                    let port_name = source.get_output_port_name(&port);
                    let schema = source
                        .get_output_schema(&port)
                        .map_err(ExecutionError::Factory)?;
                    create_edge(
                        &mut edges,
                        edge,
                        EdgeKind::FromSource {
                            port_type,
                            port_name,
                        },
                        schema,
                    );
                }
            }

            NodeKind::Processor(processor) => {
                let input_schemas =
                    validate_input_schemas(&dag, &edges, node_index, processor.get_input_ports())?;

                for edge in dag.graph().edges(node_index) {
                    let schema = processor
                        .get_output_schema(&edge.weight().from, &input_schemas)
                        .await
                        .map_err(ExecutionError::Factory)?;
                    create_edge(&mut edges, edge, EdgeKind::FromProcessor, schema);
                }
            }

            NodeKind::Sink(sink) => {
                let input_schemas =
                    validate_input_schemas(&dag, &edges, node_index, sink.get_input_ports())?;
                sink.prepare(input_schemas)
                    .map_err(ExecutionError::Factory)?;
            }
        }
    }

    Ok(dag.map_owned(
        |_, node| node,
        |edge, _| edges[edge.index()].take().expect("We traversed every edge"),
    ))
}

fn find_output_port_type(
    ports: &[OutputPortDef],
    edge: EdgeReference<DagEdgeType>,
) -> OutputPortType {
    let handle = edge.weight().from;
    for port in ports {
        if port.handle == handle {
            return port.typ;
        }
    }
    panic!("BUG: port {handle} not found")
}

fn create_edge(
    edges: &mut [Option<EdgeType>],
    edge: EdgeReference<DagEdgeType>,
    edge_kind: EdgeKind,
    schema: Schema,
) {
    let edge_ref = &mut edges[edge.id().index()];
    debug_assert!(edge_ref.is_none());
    *edge_ref = Some(EdgeType::new(
        edge.weight().from,
        edge.weight().to,
        schema,
        edge_kind,
    ));
}

fn validate_input_schemas(
    dag: &daggy::Dag<NodeType, DagEdgeType>,
    edge_and_contexts: &[Option<EdgeType>],
    node_index: NodeIndex,
    input_ports: Vec<PortHandle>,
) -> Result<HashMap<PortHandle, Schema>, ExecutionError> {
    let node_handle = &dag.graph()[node_index].handle;

    let mut input_schemas = HashMap::new();
    for edge in dag.graph().edges_directed(node_index, Direction::Incoming) {
        let port_handle = edge.weight().to;

        let edge = edge_and_contexts[edge.id().index()].as_ref().expect(
            "This edge has been created from the source node because we traverse in topological order"
        );

        if input_schemas
            .insert(port_handle, edge.schema.clone())
            .is_some()
        {
            return Err(ExecutionError::DuplicateInput {
                node: node_handle.clone(),
                port: port_handle,
            });
        }
    }

    for port in input_ports {
        if !input_schemas.contains_key(&port) {
            return Err(ExecutionError::MissingInput {
                node: node_handle.clone(),
                port,
            });
        }
    }
    Ok(input_schemas)
}

#[cfg(test)]
mod tests {
    use dozer_types::node::NodeHandle;

    use super::*;

    use crate::{
        tests::{
            processors::{ConnectivityTestProcessorFactory, NoInputPortProcessorFactory},
            sinks::{ConnectivityTestSinkFactory, NoInputPortSinkFactory},
            sources::ConnectivityTestSourceFactory,
        },
        DEFAULT_PORT_HANDLE,
    };

    #[test]
    #[should_panic]
    fn source_with_no_outgoing_edge_should_panic() {
        let mut dag = Dag::new();
        dag.add_source(
            NodeHandle::new(None, "source".to_string()),
            Box::new(ConnectivityTestSourceFactory),
        );
        validate_connectivity(&dag);
    }

    #[test]
    #[should_panic]
    fn processor_with_no_outgoing_edge_should_panic() {
        let mut dag = Dag::new();
        let source = dag.add_source(
            NodeHandle::new(None, "source".to_string()),
            Box::new(ConnectivityTestSourceFactory),
        );
        let processor = dag.add_processor(
            NodeHandle::new(None, "processor".to_string()),
            Box::new(ConnectivityTestProcessorFactory),
        );
        dag.connect_with_index(source, DEFAULT_PORT_HANDLE, processor, DEFAULT_PORT_HANDLE)
            .unwrap();
        validate_connectivity(&dag);
    }

    #[test]
    #[should_panic]
    fn sink_with_no_input_port_should_panic() {
        let mut dag = Dag::new();
        dag.add_sink(
            NodeHandle::new(None, "sink".to_string()),
            Box::new(NoInputPortSinkFactory),
        );
        validate_connectivity(&dag);
    }

    #[test]
    #[should_panic]
    fn processor_with_no_input_port_should_panic() {
        let mut dag = Dag::new();
        let processor = dag.add_processor(
            NodeHandle::new(None, "processor".to_string()),
            Box::new(NoInputPortProcessorFactory),
        );
        let sink = dag.add_sink(
            NodeHandle::new(None, "sink".to_string()),
            Box::new(ConnectivityTestSinkFactory),
        );
        dag.connect_with_index(processor, DEFAULT_PORT_HANDLE, sink, DEFAULT_PORT_HANDLE)
            .unwrap();
        validate_connectivity(&dag);
    }

    #[test]
    #[should_panic]
    fn sink_with_unconnected_input_port_should_panic() {
        let mut dag = Dag::new();
        dag.add_sink(
            NodeHandle::new(None, "sink".to_string()),
            Box::new(ConnectivityTestSinkFactory),
        );
        validate_connectivity(&dag);
    }

    #[test]
    #[should_panic]
    fn sink_with_over_connected_input_port_should_panic() {
        let mut dag = Dag::new();
        let source1 = dag.add_source(
            NodeHandle::new(None, "source1".to_string()),
            Box::new(ConnectivityTestSourceFactory),
        );
        let source2 = dag.add_source(
            NodeHandle::new(None, "source2".to_string()),
            Box::new(ConnectivityTestSourceFactory),
        );
        let sink = dag.add_sink(
            NodeHandle::new(None, "sink".to_string()),
            Box::new(ConnectivityTestSinkFactory),
        );
        dag.connect_with_index(source1, DEFAULT_PORT_HANDLE, sink, DEFAULT_PORT_HANDLE)
            .unwrap();
        dag.connect_with_index(source2, DEFAULT_PORT_HANDLE, sink, DEFAULT_PORT_HANDLE)
            .unwrap();
        validate_connectivity(&dag);
    }

    #[test]
    #[should_panic]
    fn processor_with_unconnected_input_port_should_panic() {
        let mut dag = Dag::new();
        let processor = dag.add_processor(
            NodeHandle::new(None, "processor".to_string()),
            Box::new(ConnectivityTestProcessorFactory),
        );
        let sink = dag.add_sink(
            NodeHandle::new(None, "sink".to_string()),
            Box::new(ConnectivityTestSinkFactory),
        );
        dag.connect_with_index(processor, DEFAULT_PORT_HANDLE, sink, DEFAULT_PORT_HANDLE)
            .unwrap();
        validate_connectivity(&dag);
    }

    #[test]
    #[should_panic]
    fn processor_with_over_connected_input_port_should_panic() {
        let mut dag = Dag::new();
        let source1 = dag.add_source(
            NodeHandle::new(None, "source1".to_string()),
            Box::new(ConnectivityTestSourceFactory),
        );
        let source2 = dag.add_source(
            NodeHandle::new(None, "source2".to_string()),
            Box::new(ConnectivityTestSourceFactory),
        );
        let processor = dag.add_processor(
            NodeHandle::new(None, "processor".to_string()),
            Box::new(ConnectivityTestProcessorFactory),
        );
        let sink = dag.add_sink(
            NodeHandle::new(None, "sink".to_string()),
            Box::new(ConnectivityTestSinkFactory),
        );
        dag.connect_with_index(source1, DEFAULT_PORT_HANDLE, processor, DEFAULT_PORT_HANDLE)
            .unwrap();
        dag.connect_with_index(source2, DEFAULT_PORT_HANDLE, processor, DEFAULT_PORT_HANDLE)
            .unwrap();
        dag.connect_with_index(processor, DEFAULT_PORT_HANDLE, sink, DEFAULT_PORT_HANDLE)
            .unwrap();
        validate_connectivity(&dag);
    }

    #[test]
    fn validate_source_sink_dag() {
        let mut dag = Dag::new();
        let source = dag.add_source(
            NodeHandle::new(None, "source".to_string()),
            Box::new(ConnectivityTestSourceFactory),
        );
        let sink = dag.add_sink(
            NodeHandle::new(None, "sink".to_string()),
            Box::new(ConnectivityTestSinkFactory),
        );
        dag.connect_with_index(source, DEFAULT_PORT_HANDLE, sink, DEFAULT_PORT_HANDLE)
            .unwrap();
        validate_connectivity(&dag);
    }

    #[test]
    fn validate_link_shaped_dag() {
        let mut dag = Dag::new();
        let source = dag.add_source(
            NodeHandle::new(None, "source".to_string()),
            Box::new(ConnectivityTestSourceFactory),
        );
        let processor = dag.add_processor(
            NodeHandle::new(None, "processor1".to_string()),
            Box::new(ConnectivityTestProcessorFactory),
        );
        let sink = dag.add_sink(
            NodeHandle::new(None, "sink".to_string()),
            Box::new(ConnectivityTestSinkFactory),
        );
        dag.connect_with_index(source, DEFAULT_PORT_HANDLE, processor, DEFAULT_PORT_HANDLE)
            .unwrap();
        dag.connect_with_index(processor, DEFAULT_PORT_HANDLE, sink, DEFAULT_PORT_HANDLE)
            .unwrap();
        validate_connectivity(&dag);
    }
}
